#include <float.h>
#include <stdio.h>
#include "maxpool_layer.h"

void forward_maxpool_layer(const maxpool_layer l, network net)
{
    int b,i,j,k,m,n;
    int w_offset = -l.pad/2;
    int h_offset = -l.pad/2;
    int h = l.out_h;
    int w = l.out_w;
    int c = l.out_c;

    for (b = 0; b < l.batch; b++) {
        for (k = 0; k < c; k++) {
            for (i = 0; i < h; i++) {
                for (j = 0; j < w; j++) {
                    int out_index = j + w * (i + h * (k + c * b));
                    float max = -FLT_MAX;
                    int max_i = -1;
                    for (n = 0; n < l.size; n++) {
                        for (m = 0; m < l.size; m++) {
                            int cur_h = h_offset + i * l.stride + n;
                            int cur_w = w_offset + j * l.stride + m;
                            int index = cur_w + l.w*(cur_h + l.h * (k + b*l.c));
                            int valide = ((cur_h >= 0) && (cur_h < l.h) &&
                                        (cur_w >= 0) && (cur_w < l.w));
                            float val = (valide != 0) ? net.input[index] : -FLT_MAX;
                            max_i = (val > max) ? index : max_i;
                            max = (val > max) ? val : max;
                        }
                    }
                    l.output[out_index] = max;
                    l.indexes[out_index] = max_i;
                }
            }
        }
    }
}
maxpool_layer make_maxpool_layer(int batch, int h, int w, int c, int size, int stride, int padding)
{
    maxpool_layer l = {0};
    l.type = MAXPOOL;
    l.batch = batch;
    l.h = h;
    l.w = w;
    l.c = c;
    l.pad = padding;
    l.out_w = (w + padding - size)/stride + 1;
    l.out_h = (h + padding - size)/stride + 1;
    l.out_c = c;
    l.outputs = l.out_h * l.out_w * l.out_c;
    l.inputs = h * w * c;
    l.size = size;
    l.stride = stride;
    int output_size = l.out_h * l.out_w * l.out_c * batch;
    l.output = calloc(output_size, sizeof(float));
    l.indexes = calloc(output_size, sizeof(int));
    l.forward = forward_maxpool_layer;

    return l;
}